
from video import VideoRequest, generate_video
from images import generate_images
from mcp.server import FastMCP
from loguru import logger


def add_jimeng_tools(mcp: FastMCP):

    JIMENG_API_TOKEN = "7433e85ddc99ed88aa8656aacf6630dc"  # 你登录即梦获得的session_id

    @mcp.tool("use_description")
    async def list_tools():
        """列出所有可用的工具及其参数"""
        return {
            "tools": [
                {
                    "name": "生成图片",
                    "description": "根据文本描述生成图片",
                    "parameters": {
                        "prompt": {
                            "type": "string",
                            "description": "图片的文本prompt描述,800字符长度限制，一个汉字算一个字符长度",
                            "required": True
                        },
                        "sample_strength": {
                            "type": "number",
                            "description": "生成图片的精细度(可选,范围0-1,默认0.5)",
                            "required": False
                        },
                        "width": {
                            "type": "number",
                            "description": "生成图片的宽度(可选,默认1024，最大1024)",
                            "required": False
                        },
                        "height": {
                            "type": "number",
                            "description": "生成图片的高度(可选,默认1024，最大1024)",
                            "required": False
                        }
                    }
                },
                {
                    "name": "生成视频",
                    "description": "根据文本描述生成视频",
                    "parameters": {
                        "prompt": {
                            "type": "string",
                            "description": "视频的文本prompt描述, 800字符长度限制，一个汉字算一个字符长度",
                            "required": True
                        },
                        "aspect_ratio": {
                            "type": "string",
                            "description": "生成视频的宽高比(默认为 16:9)",
                            "required": False
                        },
                        "duration_ms": {
                            "type": "number",
                            "description": "视频时长（毫秒），默认为 5000",
                            "required": False
                        },
                        "fps": {
                            "type": "number",
                            "description": "视频帧率，默认为 24",
                            "required": False
                        }
                    }
                }
            ]
        }

    @mcp.tool("generate_image")
    async def generate_image(
            prompt: str,
            sample_strength: float = 0.5,
            width: int = 1024,
            height: int = 1024
    ):
        """根据文本描述生成图片
        Args:
            prompt: 图片的文本prompt描述，800字符长度限制，一个汉字算一个字符长度
            sample_strength: 生成图片的精细度(可选,范围0-1,默认0.5)
            width: 生成图片的宽度(可选,默认1024，最大1024)
            height: 生成图片的高度(可选,默认1024，最大1024)

        Returns:
            List: 包含生成结果的JSON字符串
        """
        logger.info(f"收到生成请求: {prompt}")

        # 验证prompt参数
        if not prompt or len(prompt) > 800:
            error_msg = "prompt不能为空,且长度不能超过800"
            logger.error(error_msg)
            return {
                "success": False,
                "error": error_msg,
                "images": []
            }

        # 验证sample_strength参数范围
        if not 0 <= sample_strength <= 1:
            error_msg = f"sample_strength参数必须在0-1范围内,当前值: {sample_strength}"
            logger.error(error_msg)
            return {
                "success": False,
                "error": error_msg,
                "images": []
            }

        # 验证width和height参数
        if width <= 0 or height <= 0 or width > 1024 or height > 1024:
            error_msg = f"width和height必须大于0,小于1024，当前值: width={width}, height={height}"
            logger.error(error_msg)
            return {
                "success": False,
                "error": error_msg,
                "images": []
            }

        try:
            image_urls = generate_images(
                model="jimeng-2.1",
                prompt=prompt,
                width=width,
                height=height,
                sample_strength=sample_strength,
                negative_prompt="",
                refresh_token=JIMENG_API_TOKEN
            )

            # 返回结果
            return {
                "success": True,
                "error": None,
                "images": image_urls
            }

        except Exception as e:
            error_msg = f"生成图片失败: {str(e)}"
            logger.error(error_msg)
            return {
                "success": False,
                "error": error_msg,
                "images": []
            }

    # 添加自定义 MCP 工具
    @mcp.tool("generate_video")
    async def generate_video_mcp(
            prompt: str,
            aspect_ratio: str = "16:9",
            duration_ms: int = 5000,
            fps: int = 24,
    ):
        """
        生成一个基于文本提示的 AI 视频。
        Args:
            prompt: 用于生成视频的文本提示词
            aspect_ratio: 视频宽高比，默认为 "16:9"
            duration_ms: 视频时长（毫秒），默认为 5000
            fps: 视频帧率，默认为 24
        Returns:
            dict: 包含以下字段的字典：
                - success: 是否成功生成
                - error: 报错信息
                - video: 生成视频的 URL
        """
        request = VideoRequest(
            prompt=prompt,
            aspect_ratio=aspect_ratio,
            duration_ms=duration_ms,
            fps=fps
        )
        res = await generate_video(request)
        return res